#pragma once

// COMPI Runtime Context
//
// 目标与分层：
// - 上下文（Contex）：统一对外入口，提供 communicator、资源访问与 MPI 调用封装
// - 通道（Channel）：数据路径，支持可选的 per-channel communicator（见 channel.h）
// - 帧协议（Framing）：显式 TLV 帧协议（见 frame.h）
// - 资源管理（Tags/Types/Requests）：每 communicator 的真实单例 TagAllocator 等（见
// resource_manager.h）
//
// 骨架：每 communicator 真实单例的 TagAllocator + 可恢复错误栈（ErrorStack）
// + 显式 TLV 帧协议 + 可选 per-channel communicator。

#include "compi/result.h"
#include "compi/tag.h"

#include <chrono>
#include <functional>
#include <memory>
#include <mpi.h>
#include <optional>
#include <string>
#include <vector>

namespace compi
{

// 预声明：Context 作为策略载体
class Context;

// 错误策略类型：定义失败行为（抛异常、记录并继续、终止进程）
enum class FailurePolicy
{
    ThrowException,
    LogAndContinue,
    AbortProcess
};

// 统一 MPI 调用返回码处理（依赖 Context 注入失败策略）：
// - 返回 true 表示 rc == MPI_SUCCESS；
// - 在 ThrowException 策略下抛异常；
// - 在 LogAndContinue 下记录并返回 false；
// - 在 AbortProcess 下调用 MPI_Abort 并终止进程。
// 注意：此函数的具体实现位于 src/compi.cpp，由 Context/Environment 提供策略。
bool call_mpi(Context & ctx,
              int rc,
              const char * operation_name,
              std::optional<FailurePolicy> policy_override = std::nullopt);

inline bool call_mpi(int rc,
                     const char * operation_name,
                     MPI_Comm comm,
                     std::optional<FailurePolicy> policy_override = std::nullopt);

// 全局失败策略（由 Environment 管理）：
// - 当无法获得 Context/Environment 时使用该策略（例如空的 TransportRequest::wait）
// - 由 Environment::set_failure_policy 写入；默认值为 ThrowException。
FailurePolicy global_failure_policy() noexcept;
void set_global_failure_policy(FailurePolicy policy) noexcept;

// ===============
// Environment API
// ===============

// 线程支持级别（与 MPI 线程级别一一对应）
enum class ThreadLevel
{
    Single,     // MPI_THREAD_SINGLE
    Funneled,   // MPI_THREAD_FUNNELED
    Serialized, // MPI_THREAD_SERIALIZED
    Multiple    // MPI_THREAD_MULTIPLE
};

// 创建并初始化环境；如果底层是 MPI，则完成 MPI_Init。
// 参数：无
// 返回值：无
// 异常：在初始化失败时根据策略选择抛出或终止。
class Environment
{
public:
    // 使用 MPI_Init_thread 初始化，默认请求 Multiple。
    Environment(ThreadLevel req = ThreadLevel::Multiple);
    ~Environment();

    // 禁止拷贝，允许移动。
    Environment(const Environment &) = delete;
    Environment & operator=(const Environment &) = delete;
    Environment(Environment &&) noexcept = default;
    Environment & operator=(Environment &&) noexcept = default;

    // 设置失败策略。
    void set_failure_policy(FailurePolicy policy) noexcept;

    // 获取失败策略。
    FailurePolicy failure_policy() const noexcept;

    // 线程级别
    ThreadLevel provided_thread_level() const noexcept { return _provided_level; }
    ThreadLevel requested_thread_level() const noexcept { return _requested_level; }

private:
    ThreadLevel _requested_level{ThreadLevel::Single};
    ThreadLevel _provided_level{ThreadLevel::Single};
    // 仅当本实例负责调用 MPI_Init_* 时才在析构中调用 MPI_Finalize。
    bool _owns_mpi{false};
    // 全局默认失败策略（由 Environment 管理）
    FailurePolicy _failure_policy{FailurePolicy::ThrowException};
};

/**
 * Context：COMPI 的轻量运行时上下文
 *
 * 设计约束：
 * - 不再承担 MPI 调用封装、日志/度量、缓冲池等冗余职责；
 * - 聚焦于：为上层提供关联的 MPI_Comm、Environment 以及每 communicator 的 TagAllocator；
 * - 失败策略与可恢复错误栈由 Context 负责；
 */

// ==================
// Context and Errors
// ==================

/**
 * Context：统一所有MPI调用的入口点与失败策略实施
 * - 每个MPI_Comm对应唯一Context实例
 */
class Context
{
public:
    // 获取或创建给定通信器的实例（不要求提供 Environment）
    static Context & for_comm(MPI_Comm comm);
    // 获取或创建给定通信器的实例，并记录 Environment 引用
    static Context & for_comm(MPI_Comm comm, Environment & env);
    // 销毁给定通信器的实例
    static void destroy_for_comm(MPI_Comm comm);

    Result<void> call_mpi(std::function<int()> mpi_fn,
                          const char * operation_name,
                          std::optional<FailurePolicy> policy_override = std::nullopt);

    template <typename T>
    Result<T> call_mpi(std::function<int()> mpi_fn,
                       const char * operation_name,
                       std::function<T()> result_extractor,
                       std::optional<FailurePolicy> policy_override = std::nullopt);

    MPI_Comm comm() const noexcept;
    int rank() const noexcept;
    int size() const noexcept;

    // 失败策略：默认由 Environment 提供；Context 可设置覆盖策略。
    void set_failure_policy(FailurePolicy policy) noexcept;
    FailurePolicy failure_policy() const noexcept;           // 返回覆盖值或默认值
    FailurePolicy effective_failure_policy() const noexcept; // 等价于 failure_policy()

    struct MpiStats
    {
        std::size_t total_calls = 0;
        std::size_t failed_calls = 0;
        std::size_t bytes_sent = 0;
        std::size_t bytes_received = 0;
        double total_time_ms = 0.0;
    };

    MpiStats stats() const;
    void reset_stats() noexcept;

    class ErrorStack & error_stack() noexcept;

    // 资源访问（合并自原 Context）
    MessageTagAllocator & tag_allocator();
    // 获取关联的 Environment；如果未通过 for_comm(comm, env) 设定，则返回错误。
    Result<Environment &> environment() const noexcept;

    // 统一处理错误（基于 Result<void>）：
    // - 根据当前失败策略选择抛出、记录并继续或终止进程；
    // - 同时将错误写入可恢复错误栈。
    void handle_error(const Result<void> & res);

public:
    explicit Context(MPI_Comm comm);
    ~Context();

private:
    bool initialize_comm_properties();
    struct Impl;
    std::unique_ptr<Impl> _impl;
};

/**
 * 可恢复错误栈：管理MPI操作中的错误状态
 */
class ErrorStack
{
public:
    struct ErrorEntry
    {
        ResultCode code;
        std::string message;
        std::string operation;
        std::chrono::system_clock::time_point timestamp;
    };

    void
    push_error(ResultCode code, const std::string & message, const std::string & operation = "");
    Result<ErrorEntry> get_last_error() const;
    const ErrorEntry * last_error() const noexcept;
    std::vector<ErrorEntry> all_errors() const;
    void clear() noexcept;
    bool has_errors() const noexcept;
    std::size_t error_count() const noexcept;

private:
    struct Impl;
    std::unique_ptr<Impl> _impl;

public:
    ErrorStack();
    ~ErrorStack();

    ErrorStack(const ErrorStack &) = delete;
    ErrorStack & operator=(const ErrorStack &) = delete;
    ErrorStack(ErrorStack &&) noexcept;
    ErrorStack & operator=(ErrorStack &&) noexcept;
};

// 模板方法实现（保持与核心文件一致）
template <typename T>
Result<T>
Context::call_mpi(std::function<int()> mpi_fn,
                  const char * operation_name,
                  std::function<T()> result_extractor,
                  std::optional<FailurePolicy> policy_override)
{
    auto result = call_mpi(mpi_fn, operation_name, policy_override);
    if (result.is_error())
    {
        return Result<T>::error(result.error_code(), result.error_message());
    }

    try
    {
        T value = result_extractor();
        return Result<T>::success(std::move(value));
    }
    catch (const std::exception & e)
    {
        return Result<T>::error(ResultCode::Unknown,
                                std::string("Result extraction failed: ") + e.what());
    }
}

// 兼容封装：保留旧形式以减少大范围改动；内部通过 Context 注入策略。
inline bool
call_mpi(int rc,
         const char * operation_name,
         MPI_Comm comm,
         std::optional<FailurePolicy> policy_override)
{
    return call_mpi(Context::for_comm(comm), rc, operation_name, policy_override);
}

} // namespace compi
